module poly_rhombus_mod
  use recon_type_mod
  use recon_math_mod
  use poly_utils_mod

  implicit none

  !! 1: i,j,x^?,y^?
  !! 2: cell
  !! 3: sub-stencil
  integer, dimension(:,:), allocatable, target :: stencil_3
  integer, dimension(:,:), allocatable, target :: stencil_5
  integer, dimension(:,:), allocatable, target :: stencil_7
    
contains

  recursive subroutine poly_rhombus_init(this, nd, sw, swx, swy, swz, xc, yc, zc, dx, dy, dz, ic, jc, kc, is, ie, js, je, ks, ke, mask, id)

    class(recon_type), intent(inout) :: this
    integer , intent(in), optional :: nd
    integer , intent(in), optional :: sw
    integer , intent(in), optional :: swx
    integer , intent(in), optional :: swy
    integer , intent(in), optional :: swz
    real(16), intent(in), optional :: xc(:)
    real(16), intent(in), optional :: yc(:)
    real(16), intent(in), optional :: zc(:)
    real(16), intent(in), optional :: dx
    real(16), intent(in), optional :: dy
    real(16), intent(in), optional :: dz
    integer , intent(in), optional :: ic
    integer , intent(in), optional :: jc
    integer , intent(in), optional :: kc
    integer , intent(in), optional :: is
    integer , intent(in), optional :: ie
    integer , intent(in), optional :: js
    integer , intent(in), optional :: je
    integer , intent(in), optional :: ks
    integer , intent(in), optional :: ke
    integer , intent(in), optional :: mask(:,:,:) ! Cell mask
    integer , intent(in), optional :: id

    real(16), allocatable :: x(:), y(:)
    integer i, j, k, isub, iCell, sub_swx, sub_swy, sub_is, sub_ie, sub_js, sub_je
    integer max_ngb, ingb, m, n
    
    if(this%id==0)then
      if(.not.allocated(stencil_3   ))allocate(stencil_3   (4,5 ))
      if(.not.allocated(stencil_5   ))allocate(stencil_5   (4,13))
      if(.not.allocated(stencil_7   ))allocate(stencil_7   (4,25))
      
                                  stencil_3(:,1) = [2,3,0,2]                             !       y^2
      stencil_3(:,2) = [1,2,1,0]; stencil_3(:,3) = [2,2,0,0]; stencil_3(:,4) = [3,2,2,0] !   x    1    x^2
                                  stencil_3(:,5) = [2,1,0,1]                             !        y
      
                                                                stencil_5(:, 1) = [3,5,0,4]                                                            !                     y^4
                                   stencil_5(:, 2) = [2,4,1,2]; stencil_5(:, 3) = [3,4,0,3]; stencil_5(:, 4) = [4,4,2,2]                               !          x y^2      y^3      x^2 y^2
      stencil_5(:, 5) = [1,3,2,0]; stencil_5(:, 6) = [2,3,1,0]; stencil_5(:, 7) = [3,3,0,0]; stencil_5(:, 8) = [4,3,3,0]; stencil_5(:, 9) = [5,3,4,0]  !  x^2     x           1       x^3          x^4
                                   stencil_5(:,10) = [2,2,1,1]; stencil_5(:,11) = [3,2,0,1]; stencil_5(:,12) = [4,2,2,1]                               !          x y         y       x^2 y
                                                                stencil_5(:,13) = [3,1,0,2]                                                            !                     y^2
                                   
                                                                                             stencil_7(:, 1) = [4,7,0,6]                                 
                                                                stencil_7(:, 2) = [3,6,1,4]; stencil_7(:, 3) = [4,6,0,4]; stencil_7(:, 4) = [5,6,2,4]                      
                                   stencil_7(:, 5) = [2,5,3,2]; stencil_7(:, 6) = [3,5,1,2]; stencil_7(:, 7) = [4,5,0,2]; stencil_7(:, 8) = [5,5,2,2]; stencil_7(:, 9) = [6,5,4,2]           
      stencil_7(:,10) = [1,4,5,0]; stencil_7(:,11) = [2,4,3,0]; stencil_7(:,12) = [3,4,1,0]; stencil_7(:,13) = [4,4,0,0]; stencil_7(:,14) = [5,4,2,0]; stencil_7(:,15) = [6,4,4,0]; stencil_7(:,16) = [7,4,6,0]
                                   stencil_7(:,17) = [2,3,3,1]; stencil_7(:,18) = [3,3,1,1]; stencil_7(:,19) = [4,3,0,1]; stencil_7(:,20) = [5,3,2,1]; stencil_7(:,21) = [6,3,4,1]           
                                                                stencil_7(:,22) = [3,2,1,3]; stencil_7(:,23) = [4,2,0,3]; stencil_7(:,24) = [5,2,2,3]                      
                                                                                             stencil_7(:,25) = [4,1,0,5]     
    endif
  
    call this%clear()
    
    this%sw      = sw
    this%swx     = sw
    this%swy     = sw
    this%swz     = 1
    this%nd      = nd
    
    this%dx = dx
    this%dy = dy
    if(present(dz))then
      this%dz = dz
    else
      this%dz = 1
    endif

    if (present(id)) this%id = id

    if(present(is))then; this%is = is; else; this%is = 1       ; endif
    if(present(ie))then; this%ie = ie; else; this%ie = this%swx; endif
    if(present(js))then; this%js = js; else; this%js = 1       ; endif
    if(present(je))then; this%je = je; else; this%je = this%swy; endif
    if(present(ks))then; this%ks = ks; else; this%ks = 1       ; endif
    if(present(ke))then; this%ke = ke; else; this%ke = 1       ; endif

    allocate(this%cell_mask(this%is:this%ie,this%js:this%je,this%ks:this%ke))
    this%cell_mask = 0

    if (present(mask)) then
      this%cell_mask = mask
    else
      this%cell_mask = 1
    end if
    
    ! Set stencil
    select case (this%sw)
    case (3)
      this%ijxy => stencil_3
    case (5)
      this%ijxy => stencil_5
    case (7)
      this%ijxy => stencil_7
    end select
    this%nc = size(this%ijxy, 2)

    ! Set cell coordinates.
    allocate(this%xc(this%nc))
    allocate(this%yc(this%nc))
    allocate(x(this%is:this%ie))
    allocate(y(this%js:this%je))
    ! Set coordinates of cells on the large stencil with origin at center.
    do i = this%is, this%ie
      x(i) = -int(this%swx / 2) + i - 1
    end do
    do j = this%js, this%je
      y(j) = -int(this%swy / 2) + j - 1
    end do
    do iCell = 1, this%nc
      this%xc(iCell) = x(this%ijxy(1,iCell))
      this%yc(iCell) = y(this%ijxy(2,iCell))
    end do

    this%initialized = .true.

  end subroutine poly_rhombus_init

  subroutine poly_rhombus_calc_recon_matrix(this, ierr)

    class(recon_type), intent(inout) :: this
    integer, intent(out) :: ierr

    real(16), allocatable, dimension(:,:) :: A, iA
    integer i, j, ic, ipt, iterm, dxn, dyn

    ierr = 0

    if (allocated(this%iA        )) deallocate(this%iA        )
    if (allocated(this%poly_r16  )) deallocate(this%poly_r16  )
    if (allocated(this%recon_mtx )) deallocate(this%recon_mtx )
    if (allocated(this%dpoly_r16 )) deallocate(this%dpoly_r16 )
    if (allocated(this%drecon_mtx)) deallocate(this%drecon_mtx)

    allocate(this%iA        (this%nc ,this%nc))
    allocate(this%poly_r16  (this%npt,this%nc))
    allocate(this%recon_mtx (this%npt,this%nc))
    allocate(this%dpoly_r16 (this%npt,this%nc,0:this%swx-1,0:this%swy-1,0:this%swz-1))
    allocate(this%drecon_mtx(this%npt,this%nc,0:this%swx-1,0:this%swy-1,0:this%swz-1))

    allocate( A(this%nc,this%nc))
    allocate(iA(this%nc,this%nc))

    do ipt = 1, this%npt
      iterm = 1
      do ic = 1, this%nc
        call calc_monomial(this%x(ipt), this%y(ipt), this%ijxy(3,ic), this%ijxy(4,ic), this%poly_r16(ipt,ic))
        do dyn = 0, this%swy - 1
          do dxn = 0, this%swx - 1
            call calc_deriv_monomial(this%x(ipt), this%y(ipt), this%ijxy(3,ic), this%ijxy(4,ic), dxn, dyn, this%sw, this%sw, this%dpoly_r16(ipt,iterm,dxn,dyn,0))
            this%dpoly_r16(ipt,iterm,dxn,dyn,0) = this%dpoly_r16(ipt,iterm,dxn,dyn,0) / ( this%dx**dxn * this%dy**dyn )
          end do
        end do
        iterm = iterm + 1
      end do
    end do
    
    call calc_poly_integral_matrix(this%nc, this%nc, this%ijxy(3:4,:), this%xc, this%yc, A)
    call inverse_matrix(A, iA, ierr)
    if (ierr /= 0) then
      deallocate(A, iA)
      return
    end if

    this%recon_mtx = matmul(this%poly_r16, iA)
    
    do dyn = 0, this%swy - 1
      do dxn = 0, this%swx - 1
        this%drecon_mtx(:,:,dxn,dyn,0) = matmul(this%dpoly_r16(:,:,dxn,dyn,0), iA)
      end do
    end do

    deallocate(A, iA)

  end subroutine poly_rhombus_calc_recon_matrix
  
  subroutine poly_rhombus_reconstruct(this, fi, fo, TCI, ierr)

    class(recon_type), intent(inout) :: this
    real(8), intent(in ) :: fi(:,:,:) ! Cell averaged function values
    real(8), intent(out) :: fo(:)     ! Reconstructed function values on evaluation points
    logical, intent(in ), optional :: TCI ! Trouble Cell Indicator(TCI), 1 for existing TC(use WENO), 0 for no TC(use poly)
    integer, intent(out), optional :: ierr
    
    integer i,j,k
    integer iCell
    
    real(8) f(this%nc)
    
#ifdef DEBUG
    if(present(ierr))then
      ierr = 0
      if (size(fi) /= this%nc) then
        if (present(ierr)) ierr = 1
        return
      end if
      if (size(fo) /= this%npt) then
        if (present(ierr)) ierr = 1
        return
      end if
    endif
#endif
    
    do iCell = 1, this%nc
      f(iCell) = fi(this%ijxy(1,iCell),this%ijxy(2,iCell),1)
    enddo
    fo = matmul(this%recon_mtx, f)
      
  end subroutine poly_rhombus_reconstruct
  
  subroutine poly_rhombus_reconstruct_deriv(this, fi, fx, fy, fz, ierr)

    class(recon_type), intent(inout) :: this
    real(8), intent(in ) :: fi(:,:,:) ! Cell averaged function values
    real(8), intent(out), optional :: fx(:)     ! Reconstructed x derivative values on evaluation points
    real(8), intent(out), optional :: fy(:)     ! Reconstructed y derivative values on evaluation points
    real(8), intent(out), optional :: fz(:)     ! Reconstructed y derivative values on evaluation points
    integer, intent(out), optional :: ierr

    integer i,j,k
    integer iCell
    
    real(8) f(this%nc)
    
#ifdef DEBUG
    if(present(ierr))then
      ierr = 0
      if (size(fi) /= this%nc) then
        if (present(ierr)) ierr = 1
        return
      end if
      if (size(fo) /= this%npt) then
        if (present(ierr)) ierr = 1
        return
      end if
    endif
#endif

    do iCell = 1, this%nc
      f(iCell) = fi(this%ijxy(1,iCell),this%ijxy(2,iCell),1)
    enddo
    fx = matmul(this%drecon_mtx(:,:,1,0,0), f)
    fy = matmul(this%drecon_mtx(:,:,0,1,0), f)
    
  end subroutine poly_rhombus_reconstruct_deriv

end module poly_rhombus_mod
